'''
Author: zhao-leo 18055219130@163.com
Date: 2024-10-22 21:42:31
LastEditTime: 2024-10-22 22:15:31
'''
import git
import toml
import os

class modelDownloader():
  def __init__(self,config_path:str,cachedir:str='./.cache/models/'):
    '''
    config_path: str
      path to toml file containing model config
    cachedir: str
      path to directory to store models
    '''
    self.config_path = config_path
    self.config = self.__load_config(config_path)
    self.cachedir = cachedir
  
  def load_models(self): # load models from config
    for model in self.config:
      url = self.config[model]['url']
      path = self.cachedir+model
      if not os.path.exists(path):
        print(f"Downloading {model} from {url} to {path}")
        self.__download_model(url,path)
      else:
        print(f"Model {model} already exists in {path}\n")

  def add_model(self,model:str,url:str)->None: # add model to config
    if model in self.config:
      print(f"Model {model} already exists in {self.config_path}")
    else:
      self.config[model] = {'url':url}
      with open(self.config_path,'w') as f:
        toml.dump(self.config,f)
      print(f"Added {model} to {self.config_path}")

  def remove_model(self,model:str)->None: # remove model from config
    if model in self.config:
      self.config.pop(model)
      with open(self.config_path,'w') as f:
        toml.dump(self.config,f)
      print(f"Removed {model} from {self.config}")
    else:
      print(f"Model {model} not found in {self.config}")
  
  def __str__(self,model:str)->str:
    return self.cachedir+model
  
  def __load_config(self,config_path:str) -> dict: # load config from toml file
    with open(config_path,'r') as f:
      config = toml.load(f)
    return config
  
  def __download_model(self,url:str,path:str): # download model from url to path
    git.Repo.clone_from(url=url,to_path=path)
    print(f"Downloaded {path}")